import torch

class BinaryNeuralSwitch(torch.nn.Module):
    def __init__(self,
        dim: int | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        if dim is not None:
            self.linear = torch.nn.Linear(dim, 1, dtype=dtype, device=device)
        else:
            self.linear = torch.nn.LazyLinear(1, dtype=dtype, device=device)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        logit = torch.squeeze(self.linear(x), dim=-1)
        prob = torch.sigmoid(logit)
        result = torch.rand_like(prob) < prob
        return result, logit
